TypeTree support in autodiff #144197

KMJ-007 · 2025-07-19T22:18:34Z

TypeTrees for Autodiff

What are TypeTrees?

Memory layout descriptors for Enzyme. Tell Enzyme exactly how types are structured in memory so it can compute derivatives efficiently.

Structure

TypeTree(Vec<Type>)

Type {
    offset: isize,  // byte offset (-1 = everywhere)
    size: usize,    // size in bytes
    kind: Kind,     // Float, Integer, Pointer, etc.
    child: TypeTree // nested structure
}

Example: `fn compute(x: &f32, data: &[f32]) -> f32`

Input 0: x: &f32

TypeTree(vec![Type {
    offset: -1, size: 8, kind: Pointer,
    child: TypeTree(vec![Type {
        offset: -1, size: 4, kind: Float,
        child: TypeTree::new()
    }])
}])

Input 1: data: &[f32]

TypeTree(vec![Type {
    offset: -1, size: 8, kind: Pointer,
    child: TypeTree(vec![Type {
        offset: -1, size: 4, kind: Float,  // -1 = all elements
        child: TypeTree::new()
    }])
}])

Output: f32

TypeTree(vec![Type {
    offset: -1, size: 4, kind: Float,
    child: TypeTree::new()
}])

Why Needed?

Enzyme can't deduce complex type layouts from LLVM IR
Prevents slow memory pattern analysis
Enables correct derivative computation for nested structures
Tells Enzyme which bytes are differentiable vs metadata

What Enzyme Does With This Information:

Without TypeTrees (current state):

; Enzyme sees generic LLVM IR:
define float @distance(ptr* %p1, ptr* %p2) {
; Has to guess what these pointers point to
; Slow analysis of all memory operations
; May miss optimization opportunities
}

With TypeTrees (our implementation):

define "enzyme_type"="{[]:Float@float}" float @distance(
    ptr "enzyme_type"="{[]:Pointer}" %p1, 
    ptr "enzyme_type"="{[]:Pointer}" %p2
) {
; Enzyme knows exact type layout
; Can generate efficient derivative code directly
}

TypeTrees - Offset and -1 Explained

Type Structure

Type {
    offset: isize, // WHERE this type starts
    size: usize,   // HOW BIG this type is
    kind: Kind,    // WHAT KIND of data (Float, Int, Pointer)
    child: TypeTree // WHAT'S INSIDE (for pointers/containers)
}

Offset Values

Regular Offset (0, 4, 8, etc.)

Specific byte position within a structure

struct Point {
    x: f32, // offset 0, size 4
    y: f32, // offset 4, size 4
    id: i32, // offset 8, size 4
}

TypeTree for &Point (internal representation):

TypeTree(vec![
    Type { offset: 0, size: 4, kind: Float },   // x at byte 0
    Type { offset: 4, size: 4, kind: Float },   // y at byte 4
    Type { offset: 8, size: 4, kind: Integer }  // id at byte 8
])

Generates LLVM:

"enzyme_type"="{[]:Float@float}"

Offset -1 (Special: "Everywhere")

Means "this pattern repeats for ALL elements"

Example 1: Array `[f32; 100]`

TypeTree(vec![Type {
    offset: -1, // ALL positions
    size: 4,    // each f32 is 4 bytes
    kind: Float, // every element is float
}])

Instead of listing 100 separate Types with offsets 0,4,8,12...396

Example 2: Slice `&[i32]`

// Pointer to slice data
TypeTree(vec![Type {
    offset: -1, size: 8, kind: Pointer,
    child: TypeTree(vec![Type {
        offset: -1, // ALL slice elements
        size: 4,    // each i32 is 4 bytes
        kind: Integer
    }])
}])

Example 3: Mixed Structure

struct Container {
    header: i64,        // offset 0
    data: [f32; 1000],  // offset 8, but elements use -1
}

TypeTree(vec![
    Type { offset: 0, size: 8, kind: Integer }, // header
    Type { offset: 8, size: 4000, kind: Pointer,
        child: TypeTree(vec![Type {
            offset: -1, size: 4, kind: Float // ALL array elements
        }])
    }
])

KMJ-007 · 2025-07-19T23:50:35Z

Currently, I have implemented only for memcpy

KMJ-007 · 2025-07-19T23:50:53Z

r? @ZuseZ4

rustbot · 2025-07-19T23:51:03Z

Some changes occurred in compiler/rustc_ast/src/expand/autodiff_attrs.rs

cc @ZuseZ4

Some changes occurred in compiler/rustc_codegen_llvm/src/builder/autodiff.rs

cc @ZuseZ4

Some changes occurred in compiler/rustc_codegen_ssa

cc @WaffleLapkin

Some changes occurred in compiler/rustc_monomorphize/src/partitioning/autodiff.rs

cc @ZuseZ4

rustbot · 2025-07-20T06:10:39Z

Some changes occurred in compiler/rustc_codegen_gcc

cc @antoyo, @GuillaumeGomez

KMJ-007 · 2025-07-21T12:01:35Z

CI is failing, fixing them!

compiler/rustc_codegen_gcc/src/builder.rs

rustbot · 2025-07-23T17:33:55Z

Some changes occurred in src/tools/enzyme

cc @ZuseZ4

rustbot · 2025-09-19T04:13:17Z

This PR was rebased onto a different master commit. Here's a range-diff highlighting what actually changed.

Rebasing is a normal part of keeping PRs up to date, so no action is needed—this note is just to help reviewers.

Signed-off-by: Karan Janthe <[email protected]>

ZuseZ4 · 2025-09-28T03:39:40Z

Thank you for cleaning up the code and upstreaming it! The tests now look good, so let's merge it.
Do you remember if there are parts of my old PR that you didn't get to add here, just so that they won't be forgotten?
I'll also do some benchmarking to see if there are any perf changes since my prototype PR, since it's been a while.
I'll add some extra handling for i8 and might remove the memcpy handling (since you already add it to all calls) in a follow-up PR, but other than that, it should now be good enough to just let users test it and refine the implementation based on bug reports.

@bors r+

bors · 2025-09-28T03:39:42Z

📌 Commit 3ba5f19 has been approved by ZuseZ4

It is now in the queue for this repository.

KMJ-007 · 2025-09-28T07:56:09Z

Thank you for cleaning up the code and upstreaming it! The tests now look good, so let's merge it. Do you remember if there are parts of my old PR that you didn't get to add here, just so that they won't be forgotten? I'll also do some benchmarking to see if there are any perf changes since my prototype PR, since it's been a while. I'll add some extra handling for i8 and might remove the memcpy handling (since you already add it to all calls) in a follow-up PR, but other than that, it should now be good enough to just let users test it and refine the implementation based on bug reports.

@bors r+

thankyou and in your old PR, you have handled memmove memset(this file: https://github.com/EnzymeAD/rust/pull/157/files#diff-158a197387680f1912d11d2963f81632218de6857a971be95eebb964d90b4736R1129)

TypeTree support in autodiff # TypeTrees for Autodiff ## What are TypeTrees? Memory layout descriptors for Enzyme. Tell Enzyme exactly how types are structured in memory so it can compute derivatives efficiently. ## Structure ```rust TypeTree(Vec<Type>) Type { offset: isize, // byte offset (-1 = everywhere) size: usize, // size in bytes kind: Kind, // Float, Integer, Pointer, etc. child: TypeTree // nested structure } ``` ## Example: `fn compute(x: &f32, data: &[f32]) -> f32` **Input 0: `x: &f32`** ```rust TypeTree(vec![Type { offset: -1, size: 8, kind: Pointer, child: TypeTree(vec![Type { offset: -1, size: 4, kind: Float, child: TypeTree::new() }]) }]) ``` **Input 1: `data: &[f32]`** ```rust TypeTree(vec![Type { offset: -1, size: 8, kind: Pointer, child: TypeTree(vec![Type { offset: -1, size: 4, kind: Float, // -1 = all elements child: TypeTree::new() }]) }]) ``` **Output: `f32`** ```rust TypeTree(vec![Type { offset: -1, size: 4, kind: Float, child: TypeTree::new() }]) ``` ## Why Needed? - Enzyme can't deduce complex type layouts from LLVM IR - Prevents slow memory pattern analysis - Enables correct derivative computation for nested structures - Tells Enzyme which bytes are differentiable vs metadata ## What Enzyme Does With This Information: Without TypeTrees (current state): ```llvm ; Enzyme sees generic LLVM IR: define float `@distance(ptr*` %p1, ptr* %p2) { ; Has to guess what these pointers point to ; Slow analysis of all memory operations ; May miss optimization opportunities } ``` With TypeTrees (our implementation): ```llvm define "enzyme_type"="{[]:Float@float}" float `@distance(` ptr "enzyme_type"="{[]:Pointer}" %p1, ptr "enzyme_type"="{[]:Pointer}" %p2 ) { ; Enzyme knows exact type layout ; Can generate efficient derivative code directly } ``` # TypeTrees - Offset and -1 Explained ## Type Structure ```rust Type { offset: isize, // WHERE this type starts size: usize, // HOW BIG this type is kind: Kind, // WHAT KIND of data (Float, Int, Pointer) child: TypeTree // WHAT'S INSIDE (for pointers/containers) } ``` ## Offset Values ### Regular Offset (0, 4, 8, etc.) **Specific byte position within a structure** ```rust struct Point { x: f32, // offset 0, size 4 y: f32, // offset 4, size 4 id: i32, // offset 8, size 4 } ``` TypeTree for `&Point` (internal representation): ```rust TypeTree(vec![ Type { offset: 0, size: 4, kind: Float }, // x at byte 0 Type { offset: 4, size: 4, kind: Float }, // y at byte 4 Type { offset: 8, size: 4, kind: Integer } // id at byte 8 ]) ``` Generates LLVM: ```llvm "enzyme_type"="{[]:Float@float}" ``` ### Offset -1 (Special: "Everywhere") **Means "this pattern repeats for ALL elements"** #### Example 1: Array `[f32; 100]` ```rust TypeTree(vec![Type { offset: -1, // ALL positions size: 4, // each f32 is 4 bytes kind: Float, // every element is float }]) ``` Instead of listing 100 separate Types with offsets `0,4,8,12...396` #### Example 2: Slice `&[i32]` ```rust // Pointer to slice data TypeTree(vec![Type { offset: -1, size: 8, kind: Pointer, child: TypeTree(vec![Type { offset: -1, // ALL slice elements size: 4, // each i32 is 4 bytes kind: Integer }]) }]) ``` #### Example 3: Mixed Structure ```rust struct Container { header: i64, // offset 0 data: [f32; 1000], // offset 8, but elements use -1 } ``` ```rust TypeTree(vec![ Type { offset: 0, size: 8, kind: Integer }, // header Type { offset: 8, size: 4000, kind: Pointer, child: TypeTree(vec![Type { offset: -1, size: 4, kind: Float // ALL array elements }]) } ]) ```

Rollup of 6 pull requests Successful merges: - #140482 (std::net: update tcp deferaccept delay type to Duration.) - #141469 (Allow `&raw [mut | const]` for union field in safe code) - #144197 (TypeTree support in autodiff) - #146675 (Allow shared access to `Exclusive<T>` when `T: Sync`) - #147113 (Reland "Add LSX accelerated implementation for source file analysis") - #147120 (Fix --extra-checks=spellcheck to prevent cargo install every time) r? `@ghost` `@rustbot` modify labels: rollup

Rollup merge of #144197 - KMJ-007:type-tree, r=ZuseZ4 TypeTree support in autodiff # TypeTrees for Autodiff ## What are TypeTrees? Memory layout descriptors for Enzyme. Tell Enzyme exactly how types are structured in memory so it can compute derivatives efficiently. ## Structure ```rust TypeTree(Vec<Type>) Type { offset: isize, // byte offset (-1 = everywhere) size: usize, // size in bytes kind: Kind, // Float, Integer, Pointer, etc. child: TypeTree // nested structure } ``` ## Example: `fn compute(x: &f32, data: &[f32]) -> f32` **Input 0: `x: &f32`** ```rust TypeTree(vec![Type { offset: -1, size: 8, kind: Pointer, child: TypeTree(vec![Type { offset: -1, size: 4, kind: Float, child: TypeTree::new() }]) }]) ``` **Input 1: `data: &[f32]`** ```rust TypeTree(vec![Type { offset: -1, size: 8, kind: Pointer, child: TypeTree(vec![Type { offset: -1, size: 4, kind: Float, // -1 = all elements child: TypeTree::new() }]) }]) ``` **Output: `f32`** ```rust TypeTree(vec![Type { offset: -1, size: 4, kind: Float, child: TypeTree::new() }]) ``` ## Why Needed? - Enzyme can't deduce complex type layouts from LLVM IR - Prevents slow memory pattern analysis - Enables correct derivative computation for nested structures - Tells Enzyme which bytes are differentiable vs metadata ## What Enzyme Does With This Information: Without TypeTrees (current state): ```llvm ; Enzyme sees generic LLVM IR: define float ``@distance(ptr*`` %p1, ptr* %p2) { ; Has to guess what these pointers point to ; Slow analysis of all memory operations ; May miss optimization opportunities } ``` With TypeTrees (our implementation): ```llvm define "enzyme_type"="{[]:Float@float}" float ``@distance(`` ptr "enzyme_type"="{[]:Pointer}" %p1, ptr "enzyme_type"="{[]:Pointer}" %p2 ) { ; Enzyme knows exact type layout ; Can generate efficient derivative code directly } ``` # TypeTrees - Offset and -1 Explained ## Type Structure ```rust Type { offset: isize, // WHERE this type starts size: usize, // HOW BIG this type is kind: Kind, // WHAT KIND of data (Float, Int, Pointer) child: TypeTree // WHAT'S INSIDE (for pointers/containers) } ``` ## Offset Values ### Regular Offset (0, 4, 8, etc.) **Specific byte position within a structure** ```rust struct Point { x: f32, // offset 0, size 4 y: f32, // offset 4, size 4 id: i32, // offset 8, size 4 } ``` TypeTree for `&Point` (internal representation): ```rust TypeTree(vec![ Type { offset: 0, size: 4, kind: Float }, // x at byte 0 Type { offset: 4, size: 4, kind: Float }, // y at byte 4 Type { offset: 8, size: 4, kind: Integer } // id at byte 8 ]) ``` Generates LLVM: ```llvm "enzyme_type"="{[]:Float@float}" ``` ### Offset -1 (Special: "Everywhere") **Means "this pattern repeats for ALL elements"** #### Example 1: Array `[f32; 100]` ```rust TypeTree(vec![Type { offset: -1, // ALL positions size: 4, // each f32 is 4 bytes kind: Float, // every element is float }]) ``` Instead of listing 100 separate Types with offsets `0,4,8,12...396` #### Example 2: Slice `&[i32]` ```rust // Pointer to slice data TypeTree(vec![Type { offset: -1, size: 8, kind: Pointer, child: TypeTree(vec![Type { offset: -1, // ALL slice elements size: 4, // each i32 is 4 bytes kind: Integer }]) }]) ``` #### Example 3: Mixed Structure ```rust struct Container { header: i64, // offset 0 data: [f32; 1000], // offset 8, but elements use -1 } ``` ```rust TypeTree(vec![ Type { offset: 0, size: 8, kind: Integer }, // header Type { offset: 8, size: 4000, kind: Pointer, child: TypeTree(vec![Type { offset: -1, size: 4, kind: Float // ALL array elements }]) } ]) ```

Rollup of 6 pull requests Successful merges: - rust-lang/rust#140482 (std::net: update tcp deferaccept delay type to Duration.) - rust-lang/rust#141469 (Allow `&raw [mut | const]` for union field in safe code) - rust-lang/rust#144197 (TypeTree support in autodiff) - rust-lang/rust#146675 (Allow shared access to `Exclusive<T>` when `T: Sync`) - rust-lang/rust#147113 (Reland "Add LSX accelerated implementation for source file analysis") - rust-lang/rust#147120 (Fix --extra-checks=spellcheck to prevent cargo install every time) r? `@ghost` `@rustbot` modify labels: rollup

rustbot added F-autodiff `#![feature(autodiff)]` S-waiting-on-author Status: This is awaiting some action (such as code changes or more information) from the author. T-compiler Relevant to the compiler team, which will review and decide on the PR/issue. labels Jul 19, 2025

KMJ-007 mentioned this pull request Jul 19, 2025

[WIP] TypeTree support in autodiff #143490

Closed

This comment has been minimized.

Sign in to view

rustbot added the A-LLVM Area: Code generation parts specific to LLVM. Both correctness bugs and optimization-related issues. label Jul 19, 2025

This comment has been minimized.

Sign in to view

rustbot assigned ZuseZ4 Jul 19, 2025

KMJ-007 marked this pull request as ready for review July 19, 2025 23:50

rustbot added S-waiting-on-review Status: Awaiting review from the assignee but also interested parties. and removed S-waiting-on-author Status: This is awaiting some action (such as code changes or more information) from the author. labels Jul 19, 2025

This comment has been minimized.

Sign in to view

antoyo reviewed Jul 21, 2025

View reviewed changes

compiler/rustc_codegen_gcc/src/builder.rs Show resolved Hide resolved

This comment has been minimized.

Sign in to view

rustbot added has-merge-commits PR has merge commits, merge with caution. S-waiting-on-author Status: This is awaiting some action (such as code changes or more information) from the author. labels Jul 23, 2025

rust-cloud-vms bot force-pushed the type-tree branch from d3f78d7 to 94aa8bc Compare July 23, 2025 17:37

This comment has been minimized.

Sign in to view

rustbot removed the has-merge-commits PR has merge commits, merge with caution. label Jul 23, 2025

KMJ-007 added 6 commits September 19, 2025 04:11

autodiff: add TypeTree support for arrays

31541fe

autodiff: slice support in typetree

be3617b

autodiff: tuple support in typetree

7c5fbfb

autodiff: struct support in typetree

574f0b9

autodiff: fixed test to be more precise for type tree checking

4f3f0f4

autodiff: recurion added for typetree

4520926

rust-cloud-vms bot force-pushed the type-tree branch from 6339a2c to 7470015 Compare September 19, 2025 04:13

This comment has been minimized.

Sign in to view

rust-cloud-vms bot force-pushed the type-tree branch 2 times, most recently from cc0b6c4 to 2f742ea Compare September 19, 2025 04:21

This comment has been minimized.

Sign in to view

rust-cloud-vms bot force-pushed the type-tree branch from 2f742ea to 0d04346 Compare September 19, 2025 05:09

This comment has been minimized.

Sign in to view

rust-cloud-vms bot force-pushed the type-tree branch from 0d04346 to 7070455 Compare September 19, 2025 05:21

This comment has been minimized.

Sign in to view

autodiff: typetree recursive depth query from enzyme with fallback

3ba5f19

Signed-off-by: Karan Janthe <[email protected]>

rust-cloud-vms bot force-pushed the type-tree branch from 7070455 to 3ba5f19 Compare September 19, 2025 05:42

bors added S-waiting-on-bors Status: Waiting on bors to run and complete tests. Bors will change the label on completion. and removed S-waiting-on-review Status: Awaiting review from the assignee but also interested parties. labels Sep 28, 2025

matthiaskrgr mentioned this pull request Sep 28, 2025

Rollup of 6 pull requests #147128

Merged

bors merged commit c29fb2e into rust-lang:master Sep 28, 2025
10 checks passed

rustbot added this to the 1.92.0 milestone Sep 28, 2025

TypeTree support in autodiff #144197

TypeTree support in autodiff #144197

Uh oh!

Conversation

KMJ-007 commented Jul 19, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

TypeTrees for Autodiff

What are TypeTrees?

Structure

Example: fn compute(x: &f32, data: &[f32]) -> f32

Why Needed?

What Enzyme Does With This Information:

TypeTrees - Offset and -1 Explained

Type Structure

Offset Values

Regular Offset (0, 4, 8, etc.)

Offset -1 (Special: "Everywhere")

Example 1: Array [f32; 100]

Example 2: Slice &[i32]

Example 3: Mixed Structure

Uh oh!

This comment has been minimized.

This comment has been minimized.

KMJ-007 commented Jul 19, 2025

Uh oh!

KMJ-007 commented Jul 19, 2025

Uh oh!

rustbot commented Jul 19, 2025

Uh oh!

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

rustbot commented Jul 20, 2025

Uh oh!

This comment has been minimized.

KMJ-007 commented Jul 21, 2025

Uh oh!

Uh oh!

rustbot commented Jul 23, 2025

Uh oh!

This comment has been minimized.

This comment has been minimized.

rustbot commented Sep 19, 2025

Uh oh!

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

ZuseZ4 commented Sep 28, 2025

Uh oh!

bors commented Sep 28, 2025

Uh oh!

KMJ-007 commented Sep 28, 2025

Uh oh!

Uh oh!

Uh oh!

KMJ-007 commented Jul 19, 2025 •

edited

Loading

Example: `fn compute(x: &f32, data: &[f32]) -> f32`

Example 1: Array `[f32; 100]`

Example 2: Slice `&[i32]`